#include "txr_algo_dlm_demo.h"
#include <QFileDialog>
#include <QTextStream>
#include <QDir>
#include <QFile>
#include <QtXml/QDomDocument>

#include <sstream>
#include <iomanip> 

#include "../tk_utility/tk_timer.hpp"
#include <QThread>


#include <opencv2\highgui.hpp>
#include <opencv2\imgproc.hpp>
#include "monotone_cubic/monotone_cubic.hpp"
#include <QSettings>

const static int mark_color[txr_algo_dlm_demo::DLM_CLASS_NUM][3] =
{
	{ 255, 0, 0 },
	{ 0, 255, 0 },
	{ 0, 0, 255 },

	{ 255, 0, 255 },
	{ 255, 255, 0 },
	{ 0, 255, 255 },

	{ 254, 67, 101 },
	{ 131, 175, 155 },
	{ 249, 205, 173 },

	{ 220, 87, 18 },
	{ 254, 67, 101 },
	{ 131, 175, 155 },

	{ 249, 205, 173 },
	{ 220, 87, 18 },
};

#define GRAY16_TO_8(x)					(x>>8)
#define GRAY8_TO_16(x)					((x<<8)+x)

#define GRAY16_LEVEL					65536
#define GRAY16_MAX						65535

void GenGrayMap(QVector<tk_point<tk_int32>>& v_node, QVector<unsigned short>& v_map)
{
	v_map.resize(GRAY16_LEVEL);

	if (!v_node.size())
	{
		for (tk_int32 i = 0; i < GRAY16_LEVEL; ++i)
		{
			v_map[i] = i;
		}
		return;
	}
	v_node.push_front(tk_point<tk_int32>(0, 0));
	v_node.push_back(tk_point<tk_int32>(255, 255));

	tk_int32 knum = v_node.size();
	QVector<tk_double> v_x, v_y;
	for (tk_int32 i = 0; i < knum; ++i)
	{
		v_x.push_back(GRAY8_TO_16(v_node[i].x));
		v_y.push_back(GRAY8_TO_16(v_node[i].y));
	}

	QVector<tk_double> v_inter_x, v_inter_y;
	v_inter_x.resize(GRAY16_LEVEL);
	v_inter_y.resize(GRAY16_LEVEL);
	for (tk_int32 i = 0; i < GRAY16_LEVEL; ++i)
	{
		v_inter_x[i] = i;
	}
	math::monotone_cubic::monotone_cubic_spline(v_x.data(), v_y.data(), knum, v_inter_x.data(), v_inter_y.data(), GRAY16_LEVEL);
	for (tk_int32 i = 0; i < GRAY16_LEVEL; ++i)
	{
		v_inter_y[i] = qBound(0, tk_int32(v_inter_y[i]), GRAY16_MAX);
		v_map[i] = v_inter_y[i];
	}
}

void LoadGrayMapSet(QString path, QVector<unsigned short>* v_maps)
{
	for (tk_int32 i = 0; i < 3; ++i)
	{
		v_maps[i].clear();
	}
	QString _path = QDir::toNativeSeparators(path);
	if (!QFile::exists(_path))
	{
		return;
	}
	QSettings ini(_path, QSettings::IniFormat);
	printf("Load Graymap: %s \n", _path.toLocal8Bit().data());
	//Same enhancement mode for all channels
	if (ini.contains("NODE/num"))
	{
		ini.beginGroup("NODE");

		tk_int32 node_num = ini.value("num", 0).toInt();
		QVector<tk_point<tk_int32>> v_node;
		v_node.resize(node_num);

		for (tk_int32 i = 0; i < node_num; ++i)
		{
			v_node[i].x = ini.value(QString("x_%1").arg(i), 0).toInt();
			v_node[i].y = ini.value(QString("y_%1").arg(i), 0).toInt();
		}
		ini.endGroup();

		GenGrayMap(v_node, v_maps[0]);
		v_maps[2] = v_maps[0];
		v_maps[1] = v_maps[0];

		return;
	}
	//Different enhancement mode for each channel
	else if (ini.contains("NODE_2/num"))
	{
		for (tk_int32 ch = 0; ch < 3; ++ch)
		{
			ini.beginGroup(QString("NODE_%1").arg(ch));

			tk_int32 node_num = ini.value("num", 0).toInt();
			QVector<tk_point<tk_int32>> v_node;
			v_node.resize(node_num);

			for (tk_int32 i = 0; i < node_num; ++i)
			{
				v_node[i].x = ini.value(QString("x_%1").arg(i), 0).toInt();
				v_node[i].y = ini.value(QString("y_%1").arg(i), 0).toInt();
			}
			ini.endGroup();
			GenGrayMap(v_node, v_maps[ch]);
		}
		for (tk_int32 i = 0; i < GRAY16_LEVEL; ++i)
		{
			v_maps[1][i] = v_maps[0][i] / 2 + v_maps[2][i] / 2;
		}
	}
}
bool ReadImage(QString file_path, QVector<unsigned short>* v_maps, QVector<float>& v_data, int& row, int& col)
{
	using namespace cv;
	Mat src = imread(file_path.toLocal8Bit().data(), IMREAD_ANYCOLOR | IMREAD_ANYDEPTH);
	if (src.empty())
	{
		return false;
	}
	if (src.channels() != 3 && src.channels() != 1)
	{
		return false;
	}
	row = src.rows;
	col = src.cols;

	v_data.resize(row * col * 3);

	int img_size = row * col;
	if (src.channels() == 3)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3w color_value = src.ptr<Vec3w>(i)[j];
					v_data[i * col + j] = float(color_value[2]) / 65535.f;
					v_data[img_size + i * col + j] = float(color_value[1]) / 65535.f;
					v_data[img_size * 2 + i * col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3b color_value = src.ptr<Vec3b>(i)[j];
					v_data[i * col + j] = float(color_value[2]) / 255.f;
					v_data[img_size + i * col + j] = float(color_value[1]) / 255.f;
					v_data[img_size * 2 + i * col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}
	if (src.channels() == 1)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<ushort, 1> color_value = src.ptr<Vec<ushort, 1>>(i)[j];
					v_data[i * col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size + i * col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size * 2 + i * col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<uchar, 1> color_value = src.ptr<Vec<uchar, 1>>(i)[j];
					v_data[i * col + j] = float(color_value[0]) / 255.f;
					v_data[img_size + i * col + j] = float(color_value[0]) / 255.f;
					v_data[img_size * 2 + i * col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}

	tk_bool is_map_mode = v_maps[0].size() && v_maps[1].size() && v_maps[2].size();
	if (is_map_mode)
	{
		for (int ch = 0; ch < 3; ++ch)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					float& val = v_data[img_size * ch + i * col + j];
					unsigned short val_16 = unsigned short(val * float(GRAY16_MAX));
					val = float(v_maps[ch][val_16]) / float(GRAY16_MAX);
				}
			}
		}
	}
	return true;
}

void MarkImage(st_detect_unit unit, QString dest_path, QVector<st_box>& result)
{
	st_img_rgb img = unit.img;

	using namespace cv;
	Mat mat;

	mat.create(img.h, img.w, CV_8UC3);
	int img_size = img.h * img.w;
	for (int i = 0; i < img.h; i++)
	{
		for (int j = 0; j < img.w; j++)
		{
			mat.at<Vec3b>(i, j)[2] = img.pp_rgbf[0][i * img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[1] = img.pp_rgbf[1][i * img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[0] = img.pp_rgbf[2][i * img.w + j] * 255.f;
		}
	}

	for (int i = 0; i < result.size(); ++i)
	{
		st_box r = result[i];

		cv::Rect rect(r.x, r.y, r.w, r.h);
		rectangle(mat, rect,
			Scalar(
				mark_color[r.obj_id][2],
				mark_color[r.obj_id][1],
				mark_color[r.obj_id][0]),
			1, 1, 0);

		std::stringstream stream;
		stream << std::fixed << std::setprecision(2) << result[i].obj_id << "p:" << result[i].prob;

		cv::Point pt(rect.x, rect.y);
		if (rect.x > mat.cols - 100)
		{
			pt.x = rect.x - 100;
		}
		if (rect.y < 20)
		{
			pt.y += rect.height + 20;
		}
		putText(mat, stream.str(), pt, 0, 0.5, cv::Scalar(0, 0, 255), 0.5);
	}

	{
		QFileInfo file_info(QDir::toNativeSeparators(dest_path));
		QDir dir;
		dir.mkpath(file_info.path());
	}

#ifndef _DEBUG
	imwrite(dest_path.toLocal8Bit().data(), mat);
#endif // _DEBUG

}

void SaveTxt(QString dest_path, int img_gd_count, int img_ng_count, int* cls_ng_box_count)
{
	{
		QFileInfo file_info(QDir::toNativeSeparators(dest_path));
		QDir dir;
		dir.mkpath(file_info.path());
	}
	QFile::remove(dest_path);
	QFile file(dest_path);
	if (file.open(QIODevice::WriteOnly))
	{
		QTextStream stream(&file);
		stream << "img_total: " << img_gd_count + img_ng_count << "\n";
		stream << "img_gd_count: " << img_gd_count << "\n";
		stream << "img_ng_count: " << img_ng_count << "\n";
		for (int i = 0; i < txr_algo_dlm_demo::DLM_CLASS_NUM; ++i)
		{
			stream << "class-[" << i << "]: " << cls_ng_box_count[i] << "\n";
		}
	}
	file.close();
}

void SaveTxt(QString info_file, QVector<QString>& v_img_file, QVector<QVector<st_box>>& vv_img_box)
{
	{
		QFileInfo file_info(QDir::toNativeSeparators(info_file));
		QDir dir;
		dir.mkpath(file_info.path());
	}
	QFile::remove(info_file);
	QFile file(info_file);
	if (file.open(QIODevice::WriteOnly))
	{
		QTextStream stream(&file);
		for (int i = 0; i < v_img_file.size(); ++i)
		{
			for (auto box : vv_img_box[i])
			{
				stream << v_img_file[i] << ",";
				stream << box.obj_id << ",";
				stream << box.prob << ",";
				stream << box.x << ",";
				stream << box.y << ",";
				stream << box.w << ",";
				stream << box.h << "\n";
			}
		}
	}
	file.close();

}


void XmlAppendText(QDomDocument& doc, QDomElement& node, QString name, QString text)
{
	QDomElement element = doc.createElement(name);
	element.appendChild(doc.createTextNode(text));
	node.appendChild(element);
}

void SaveXml(QString dest_path, txr_algo_dlm_demo::st_img_info& info, QString* p_obj_name, QVector<st_box>& result)
{
	{
		QFileInfo file_info(QDir::toNativeSeparators(dest_path));
		QDir dir;
		dir.mkpath(file_info.path());
	}
	QFile::remove(dest_path);
	QFile file(dest_path);
	if (!file.open(QFileDevice::ReadWrite | QFileDevice::Truncate))
	{
		return;
	}
	QDomDocument doc;
	// 	QDomProcessingInstruction instruction;
	// 	instruction = doc.createProcessingInstruction("xml", "version=\"1.0\" encoding=\"UTF-8\"");
	// 	doc.appendChild(instruction);
		//====================================================
	QDomElement annotation = doc.createElement("annotation");
	doc.appendChild(annotation);

	XmlAppendText(doc, annotation, "folder", info.folder);
	XmlAppendText(doc, annotation, "filename", info.file_name);
	XmlAppendText(doc, annotation, "path", info.file_path);

	QDomElement size = doc.createElement("size");
	XmlAppendText(doc, size, "width", QString::number(info.col));
	XmlAppendText(doc, size, "height", QString::number(info.row));
	XmlAppendText(doc, size, "depth", QString::number(info.depth));
	annotation.appendChild(size);

	for (auto obj : result)
	{
		QDomElement object = doc.createElement("object");

		XmlAppendText(doc, object, "name", p_obj_name[obj.obj_id]);
		XmlAppendText(doc, object, "difficult", "0");

		QDomElement bndbox = doc.createElement("bndbox");
		XmlAppendText(doc, bndbox, "xmin", QString::number(obj.x));
		XmlAppendText(doc, bndbox, "ymin", QString::number(obj.y));
		XmlAppendText(doc, bndbox, "xmax", QString::number(obj.x + obj.w));
		XmlAppendText(doc, bndbox, "ymax", QString::number(obj.y + obj.h));
		object.appendChild(bndbox);

		annotation.appendChild(object);
	}
	//====================================================
	QTextStream stream(&file);
	doc.save(stream, 4);
	file.close();
}

txr_algo_dlm_demo::txr_algo_dlm_demo(QWidget* parent)
	: QMainWindow(parent)
{
	ui.setupUi(this);

	connect(ui.pushButtonOpenTestDir, SIGNAL(clicked()), this, SLOT(OnOpenTest()));
	connect(ui.pushButtonOpenModelDir, SIGNAL(clicked()), this, SLOT(OnOpenModel()));
	connect(ui.pushButtonTest, SIGNAL(clicked()), this, SLOT(OnTest()));

	QVBoxLayout* layout_vbox = new QVBoxLayout();
	QString obj_name[DLM_CLASS_NUM] = {
	"foreign","ceshikuai",
	};
	for (int i = 0; i < DLM_CLASS_NUM; ++i)
	{
		m_ui_obj_cfg[i] = new st_ui_obj_cfg();
		m_ui_obj_cfg[i]->checkbox_enable.setChecked(i <= 4 ? true : false);
		m_ui_obj_cfg[i]->checkbox_enable.setText(QString("class-%1").arg(i));
		m_ui_obj_cfg[i]->label_name.setText("obj_name");
		m_ui_obj_cfg[i]->spinbox_th.setValue(30);
		m_ui_obj_cfg[i]->lineedit_name.setText(obj_name[i].isEmpty() ? "null" : obj_name[i]);
		m_ui_obj_cfg[i]->grid_layout.addWidget(&m_ui_obj_cfg[i]->checkbox_enable, 0, 0);
		m_ui_obj_cfg[i]->grid_layout.addWidget(&m_ui_obj_cfg[i]->label_name, 0, 1);
		m_ui_obj_cfg[i]->grid_layout.addWidget(&m_ui_obj_cfg[i]->lineedit_name, 0, 2);
		m_ui_obj_cfg[i]->grid_layout.addWidget(&m_ui_obj_cfg[i]->spinbox_th, 0, 3);

		layout_vbox->addLayout(&m_ui_obj_cfg[i]->grid_layout);
	}
	ui.verticalLayout->addLayout(layout_vbox);

	ui.progressBar->setRange(0, 100);
	ui.progressBar->setValue(0);
}

void txr_algo_dlm_demo::OnOpenTest()
{
	QString dir = QFileDialog::getExistingDirectory(this);
	if (dir.isEmpty())
	{
		return;
	}
	ui.labelDirTest->setText(dir);
}

void txr_algo_dlm_demo::OnOpenModel()
{
	QString dir_path = QFileDialog::getExistingDirectory(this);
	if (dir_path.isEmpty())
	{
		return;
	}
	ui.labelDirModel->setText(dir_path);

	QDir dir(dir_path);
	QStringList filters;
	filters << "*.dats" << "*.fdats" << "*.gdats" << "*.json" << "*.hdats" << "*.ohdats";

	QStringList files = dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
	ui.comboBoxModel->clear();
	ui.comboBoxModel->addItems(files);
}

void txr_algo_dlm_demo::OnTest()
{
	st_test_info info;

	info.test_dir = ui.labelDirTest->text();

	info.model_path = ui.labelDirModel->text() + "\\" + ui.comboBoxModel->currentText();

	info.save_ng = ui.checkBoxSaveNg->isChecked();
	info.save_gd = ui.checkBoxSaveGd->isChecked();
	info.save_txt = ui.checkBoxSaveTxt->isChecked();
	info.save_xml = ui.checkBoxSaveXml->isChecked();

	for (int i = 0; i < DLM_CLASS_NUM; ++i)
	{
		info.class_en[i] = m_ui_obj_cfg[i]->checkbox_enable.isChecked();
		info.class_th[i] = m_ui_obj_cfg[i]->spinbox_th.value();
		info.class_name[i] = m_ui_obj_cfg[i]->lineedit_name.text();
	}

	DlmTest(info);
}

bool txr_algo_dlm_demo::LoadDlm(txr_algo_dlm_demo::st_test_info& info, TxrAlgoDlm& dlm)
{
	st_dlm_cfg cfg;

	QFileInfo model_file = QFileInfo(info.model_path);
	QString model_name = model_file.absoluteDir().path() + "\\" + model_file.baseName();
	QString model_suffix = model_file.suffix();

	LoadGrayMapSet(model_name + ".map", m_v_gray_map);

	sprintf_s(cfg.file_path_name, MAX_PATH_BYTES, "%s", model_name.toLocal8Bit().data());
	if (model_suffix == "dats")
	{
		cfg.type = DLM_TYPE_YOLOV4_TINY;
	}
	else if (model_suffix == "fdats")
	{
		cfg.type = DLM_TYPE_YOLOV4_TINY_TRT;
	}
	else if (model_suffix == "gdats")
	{
		cfg.type = DLM_TYPE_YOLOV6_TRT;
	}
	else if (model_suffix == "hdats")
	{
		cfg.type = DLM_TYPE_YOLOV6SP5_TRT;
	}
	else if (model_suffix == "ohdats")
	{
		cfg.type = DLM_TYPE_YOLOV6SP5_OPENVINO;
	}
	return dlm.Load(cfg);
}

void txr_algo_dlm_demo::DlmTest(st_test_info& info)
{
	using namespace txr_algo_dlm;
	TxrAlgoDlm dlm;

	if (!LoadDlm(info, dlm))
	{
		return;
	}

	int batch_size = dlm.BatchSize();

	QString root = info.test_dir;
	QDir dir(root);
	QStringList folders = dir.entryList(QDir::NoDot | QDir::NoDotDot | QDir::Dirs);
	folders.append("");

	for (auto folder : folders)
	{
		ui.progressBar->setValue(0);

		info.test_dir = root + "\\" + folder;
		{
			QDir dir(info.test_dir + "\\result");
			dir.removeRecursively();
		}
		QStringList test_files;
		{
			QDir dir(info.test_dir);
			QStringList filters;
			filters << "*.png" << "*.bmp" << "*.jpg";
			test_files =
				dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
		}

		int cls_ng_box_count[DLM_CLASS_NUM] = { 0 };
		int img_gd_count = 0;
		int img_ng_count = 0;
		int file_count = 0;
		
		QVector<QString> v_img_file;
		QVector<QVector<st_box>> vv_img_box;

		QVector<QString> v_img_path(batch_size);
		QVector<QVector<float>> vv_data(batch_size);
		QVector<st_detect_unit> v_unit(batch_size);

		int batch_count = 0;
		for each (QString file in test_files)
		{
			if (batch_count < batch_size)
			{
				QString& img_path = v_img_path[batch_count];
				img_path = info.test_dir + "\\" + file;
				QVector<float>& v_data = vv_data[batch_count];
				int row, col;
				if (!ReadImage(img_path, m_v_gray_map, v_data, row, col))
				{
					continue;
				}
				st_detect_unit& unit = v_unit[batch_count];
				unit.img.w = col;
				unit.img.h = row;
				unit.img.c = 3;
				unit.img.pp_rgbf[0] = v_data.data() + col * row * 0;
				unit.img.pp_rgbf[1] = v_data.data() + col * row * 1;
				unit.img.pp_rgbf[2] = v_data.data() + col * row * 2;
				unit.param.prob_thres = 0.1;
				batch_count++;
			}
			if (batch_count != batch_size && file != test_files.last())
			{
				continue;
			}
			{
				tk::timer::universal_short_timer timer;
				dlm.Detect(v_unit.data(), batch_count);
				printf("dlm detect use time: %f ms \n", timer.elapsed_ms());
			}
			for (int i = 0; i < batch_count; ++i)
			{
				st_detect_unit& unit = v_unit[i];

				st_box* p_box = unit.result.boxes;
				tk_int32 box_size = unit.result.vaild_num;
				bool is_ng_flag = false;
				QVector<st_box> v_result_box;
				for (int i = 0; i < box_size; ++i)
				{
					int cls = p_box[i].obj_id;
					float prob = p_box[i].prob;
					if (info.class_en[cls] && prob >= float(info.class_th[cls]) / 100.f)
					{
						is_ng_flag = true;
						cls_ng_box_count[cls]++;
						v_result_box.push_back(p_box[i]);
					}
				}
				is_ng_flag ? img_ng_count += 1 : img_gd_count += 1;

				QFileInfo file_info = QFileInfo(v_img_path[i]);
				if (info.save_gd && !is_ng_flag)
				{
					QString mark_file = info.test_dir + QString("\\result\\gd\\") + file_info.baseName() + ".bmp";
					MarkImage(unit, mark_file, v_result_box);
				}
				if (info.save_ng && is_ng_flag)
				{
					QString mark_file = info.test_dir + QString("\\result\\ng\\") + file_info.baseName() + ".bmp";
					MarkImage(unit, mark_file, v_result_box);
				}

				if (info.save_xml && is_ng_flag)
				{
					QString mark_file = info.test_dir + QString("\\result\\xml\\") + file_info.baseName() + ".xml";
					st_img_info img_info;
					img_info.folder = info.test_dir;
					img_info.file_path = file_info.absoluteFilePath();
					img_info.file_name = file_info.baseName();
					img_info.row = unit.img.h;
					img_info.col = unit.img.w;
					img_info.depth = 3;
					SaveXml(mark_file, img_info, info.class_name, v_result_box);
				}

				v_img_file.push_back(file_info.baseName());
				vv_img_box.push_back(v_result_box);

				file_count = img_gd_count + img_ng_count;
				ui.progressBar->setValue(file_count * 100 / test_files.size());

				printf("total count : %d good count: %d  ng count: %d \n", file_count, img_gd_count, img_ng_count);
				printf("class: ");
				for (int i = 0; i < DLM_CLASS_NUM; ++i)
				{
					printf("[%d] = %d ,", i, cls_ng_box_count[i]);
				}
				printf("\n");

				QCoreApplication::processEvents();
			}
			batch_count = 0;
		}
		if (info.save_txt)
		{
			QString result_file = info.test_dir + QString("\\result\\") + "simulation.txt";
			SaveTxt(result_file, img_gd_count, img_ng_count, cls_ng_box_count);

			QString info_file = info.test_dir + QString("\\result\\") + "info.csv";
			SaveTxt(info_file, v_img_file, vv_img_box);

		}
		ui.progressBar->setValue(100);
	}
	printf("done!==========================================================\n");
}
